# FMCW雷达多参数估计数据集类
import os
import torch
from torch.utils.data import Dataset
from apps.fmcw.conf.app_config import AppConfig as AF

class FmcwMpeDss(Dataset):

    def __init__(self, dss_fn:str = './work/datasets/mpe_train.txt'):
        self.name = 'dss.fmcw_mpe_dss.FmcwMpeDss'
        self.Xs, self.ys = [], []
        with open(dss_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                self.Xs.append(arrs0[0])
                self.ys.append(arrs0[1])

    def __len__(self):
        return len(self.Xs)

    def __getitem__(self, idx):
        X_fn, y_fn = self.Xs[idx], self.ys[idx]
        X = torch.load(X_fn, weights_only=True)
        # print(f'X0: {X0.dtype}; {X0.shape};')
        # # X_raw = torch.load(X_fn).reshape(-1) # 类型为复数
        # print(f'X_raw: {X_raw.dtype}; {X_raw.shape};')
        # X = torch.hstack((torch.real(X_raw), torch.imag(X_raw)))
        # print(f'X: {X.dtype}; {X.shape};')
        y = torch.load(y_fn, weights_only=True)
        y = (y-AF.Y_MIN)/(AF.Y_MAX - AF.Y_MIN)
        return X.permute(2,0,1), y
    
    @staticmethod
    def generate_dss_desc() -> None:
        dss_fn = './work/datasets/mpe_train.txt'
        raw_fn = '/home/psdz/diskc/yantao/adev/iragent/work/datasets/mpe'
        fns_dict = {}
        with open(dss_fn, 'w', encoding='utf-8') as wfd:
            for root, dirs, files in os.walk(raw_fn):
                for rfn in files:
                    key = rfn[2:-3]
                    if 'X_' in rfn:
                        print(f'样本数据集文件：{rfn}; key: {rfn[2:-4]}; ffn: {root}/{rfn};')
                        fns_dict[key] = f'{root}/{rfn}'
            for root, dirs, files in os.walk(raw_fn):
                for rfn in files:
                    key = rfn[2:-3]
                    if 'y_' in rfn:
                        print(f'写入文件：{root}/{rfn};')
                        X_fn = fns_dict[key]
                        wfd.write(f'{X_fn},{root}/{rfn}\n')
